/* Copyright (c) 2025 Huawei Technologies Co., Ltd.
 * This file is a part of the CANN Open Software.
 * Licensed under CANN Open Software License Agreement Version 1.0 (the
 * "License"). Please refer to the License for details. You may not use this
 * file except in compliance with the License. THIS SOFTWARE IS PROVIDED ON AN
 * "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED,
 * INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS
 * FOR A PARTICULAR PURPOSE. See LICENSE in the root of the software repository
 * for the full text of the License.
 * ===================================================================================================================*/

#ifndef ACl_REC_H
#define ACl_REC_H
#include <memory>
#include <unordered_map>
#include <vector>

#include "acl/acl_rt.h"
#include "common.h"
#include "exe_graph/runtime/tensor.h"
namespace aclrec {
class ModelInstance {
 public:
  Result ExecuteSync();

  /**
   * 将模型执行放到指定的 stream
   * 上，该接口是异步接口，接口返回仅代表模型被下发到指定流上， 用户需要调用
   * aclrtSynchronizeStream 同步以保证模型执行完成
   */
  Result ExecuteWithStream(aclrtStream stream);

  /**
   * 如果模型满足如下特点，推荐使用该接口设置模型输入内容：
   *
   * * 静态 shape 模型
   * * 从使用场景来看，模型的所有输入都来自 Host
   * * 模型的输入数量多、且单个输入 size 小
   *
   * 该接口设置会使能批量 H-D 拷贝能力，聚合输入 Tensor 数据，使用一条拷贝指令，
   * 批量将 Host Tensor 数据拷贝到 Device，大幅提升输入数据的拷贝性能，
   * 尤为适用特征多、单项小的推荐网络
   *
   * 注意：该接口与 SetDeviceInputAddress 互斥，
   * 您需要选择一种方式管理输入内存：自己管理所有输入的 Device 地址，
   * 或者使用 GetPreAllocatedHostInput 管理所有 Host、Device 输入地址
   */
  HostAddr GetPreAllocatedHostInput(int64_t input_index) const;

  Result SetDeviceInputAddress(int64_t input_index, DevAddr dev_addr);
  Result SetDeviceOutputAddress(int64_t output_index, DevAddr dev_addr);
};

class Model;
class WeightUpdater {
 public:
  explicit WeightUpdater(Model &model);
  /**
   * 一次配置所有可更新权重，weights[i] 代表 index 为 i 的权重。
   * 权重地址同时支持 Host 与 Device 两种，
   * 通过 gert::Tensor::SetPlacement 配置权重 Tensor 位置。
   *
   * 如果您使用 aclgrphConvertToWeightRefreshableGraphs 完成的接口拆分，
   * 那么可以在对应 weights 位置传入空 Tensor（shape 为
   * [0]），代表不更新该权重。
   */
  Result SetAllWeights(const std::vector<gert::Tensor> &weights);

  Result Update();

 private:
  Model &model_;
  std::vector<gert::Tensor> weights_;
};

class Model {
 public:
  ~Model();
  static std::unique_ptr<Model> LoadFromFile(const std::string &filename);

  /**
   * 判断本模型是饭否是静态 shape 模型，
   * 对于开启了动态多档位的模型，也被认为是静态 shape 模型
   * @return
   */
  bool IsStaticShape() const;

  /**
   * 从本模型加载一个可执行实例，一个模型可以加载多分可执行实例，多个可执行实例之间可并行执行
   * @return 可执行的模型实例
   */
  std::unique_ptr<ModelInstance> LoadIns();

  /**
   * 批量更新所有权重
   * @param weights 被更新的权重，索引和数量应该与
   *     `GetAllWeightNamesToIndex`获取的结果匹配
   * @return
   */
  Result BatchUpdateWeights(const std::vector<gert::Tensor> &weights);

  /**
   * 获取权重更新器，通过权重更新器，可以更新权重，相对于`BatchUpdateWeights`函数，
   * `WeightUpdater`提供颗粒度更细的更新方式。注意：
   * **对于不支持更新权重的模型，调用该接口会导致抛出异常**
   * @return
   */
  WeightUpdater &GetWeightUpdater();

  /**
   * 根据权重节点名，获取该权重对应的`index`，
   * @param weight_names_to_index
   * @return
   */
  Result GetAllWeightNamesToIndex(std::unordered_map<std::string, int64_t> &weight_names_to_index);

 private:
  Model();
  WeightUpdater weight_updater_;
  uint32_t bundle_model_id_{kInvalidModelId};
  uint32_t update_model_id_{kInvalidModelId};
};
}  // aclrec

#endif //ACl_REC_H
